import os
import shutil

from tqdm import tqdm

sets = ['train', 'test', 'val']
# 这里使用要改成自己的类别
classes = ['drone']

# 配置文件路径
config = {
    'imagesets_dir': r'E:\pycharmProjs\yolo-ladcs\dataset\VOCdevkit-TIB\VOC2007\ImageSets\Main',
    'jpegimages_dir': r'E:\pycharmProjs\yolo-ladcs\dataset\VOCdevkit-TIB\VOC2007\JPEGImages',
    'mixed_labels_dir': r'E:\pycharmProjs\yolo-ladcs\dataset\TIB-NET-YOLO\mixed_labels',
    'labels_dir': r'E:\pycharmProjs\yolo-ladcs\dataset\TIB-NET-YOLO\labels',
    'imgs_dir': r'E:\pycharmProjs\yolo-ladcs\dataset\TIB-NET-YOLO\images',
    'YOLO-ROOT': r'E:\pycharmProjs\yolo-ladcs\dataset\TIB-NET-YOLO'
}

if __name__ == '__main__':
    os.makedirs(config['labels_dir'], exist_ok=True)
    for image_set in sets:
        tgt_labels_dir = os.path.join(config['labels_dir'], image_set)
        os.makedirs(tgt_labels_dir, exist_ok=True)
        # 读取子集图片列表
        with open(fr"{config['imagesets_dir']}\{image_set}.txt", 'r', encoding='utf-8') as image_ids_file:
            image_ids = image_ids_file.read().strip().split()
        # 移动子集图片
        for image_id in tqdm(image_ids):
            src_img_file_path = fr"{config['mixed_labels_dir']}\TIB-{image_id}.txt"
            tgt_img_file_path = fr"{tgt_labels_dir}\TIB-{image_id}.txt"
            shutil.copy(src_img_file_path, tgt_img_file_path)
